package main

import (
	"flag"
	"io"
	"log"
	"net"
	"os"
)

var (
	localPort  string
	remoteHost string
	remotePort string
)

func init() {
	flag.StringVar(&localPort, "l", "", "本地监听端口")
	flag.StringVar(&remoteHost, "h", "", "远程服务器地址")
	flag.StringVar(&remotePort, "p", "", "远程服务器端口")
	flag.Parse()

	if localPort == "" || remoteHost == "" || remotePort == "" {
		flag.Usage()
		os.Exit(1)
	}
}

func main() {
	// 启动本地监听
	listener, err := net.Listen("tcp", ":"+localPort)
	if err != nil {
		log.Fatalf("无法启动监听: %v", err)
	}
	defer listener.Close()
	log.Printf("端口转发已启动 [::]:%s -> %s:%s", localPort, remoteHost, remotePort)

	for {
		// 接受客户端连接
		clientConn, err := listener.Accept()
		if err != nil {
			log.Printf("接受连接失败: %v", err)
			continue
		}
		log.Printf("客户端连接: %s", clientConn.RemoteAddr())

		// 处理连接
		go handleConnection(clientConn)
	}
}

func handleConnection(clientConn net.Conn) {
	defer clientConn.Close()

	// 连接远程服务器
	remoteAddr := net.JoinHostPort(remoteHost, remotePort)
	serverConn, err := net.Dial("tcp", remoteAddr)
	if err != nil {
		log.Printf("连接远程服务器失败: %v", err)
		return
	}
	defer serverConn.Close()
	log.Printf("已连接到远程服务器: %s", remoteAddr)

	// 双向转发数据
	done := make(chan struct{})

	go func() {
		io.Copy(serverConn, clientConn)
		done <- struct{}{}
		// create thread
	}()

	go func() {
		io.Copy(clientConn, serverConn)
		done <- struct{}{}
	}()

	// 等待任一方向传输完成
	<-done
	log.Printf("连接关闭: %s", clientConn.RemoteAddr())
}
